import re
from urllib.parse import urlsplit
import requests
from connections import mongo_db
from loguru import logger

gt_task_col = mongo_db["annotated_ground_truth_task"]


KUBEFLOW_ENDPOINT = "http://172.25.126.219:32322"
KUBEFLOW_USERNAME = "user@example.com"
KUBEFLOW_PASSWORD = "12341234"
KUBEFLOW_API_V2 = f"{KUBEFLOW_ENDPOINT}/pipeline/apis/v2beta1/runs"


def get_istio_auth_session(url: str, username: str, password: str) -> dict:
    """
    Determine if the specified URL is secured by Dex and try to obtain a session cookie.
    WARNING: only Dex `staticPasswords` and `LDAP` authentication are currently supported
             (we default default to using `staticPasswords` if both are enabled)

    :param url: Kubeflow server URL, including protocol
    :param username: Dex `staticPasswords` or `LDAP` username
    :param password: Dex `staticPasswords` or `LDAP` password
    :return: auth session information
    """
    # define the default return object
    auth_session = {
        "endpoint_url": url,  # KF endpoint URL
        "redirect_url": None,  # KF redirect URL, if applicable
        "dex_login_url": None,  # Dex login URL (for POST of credentials)
        "is_secured": None,  # True if KF endpoint is secured
        "session_cookie": None,  # Resulting session cookies in the form "key1=value1; key2=value2"
    }

    # use a persistent session (for cookies)
    with requests.Session() as s:
        ################
        # Determine if Endpoint is Secured
        ################
        resp = s.get(url, allow_redirects=True)
        if resp.status_code != 200:
            raise RuntimeError(
                f"HTTP status code '{resp.status_code}' for GET against: {url}"
            )

        auth_session["redirect_url"] = resp.url

        # if we were NOT redirected, then the endpoint is UNSECURED
        if len(resp.history) == 0:
            auth_session["is_secured"] = False
            return auth_session
        else:
            auth_session["is_secured"] = True

        ################
        # Get Dex Login URL
        ################
        redirect_url_obj = urlsplit(auth_session["redirect_url"])

        # if we are at `/auth?=xxxx` path, we need to select an auth type
        if re.search(r"/auth$", redirect_url_obj.path):
            #######
            # TIP: choose the default auth type by including ONE of the following
            #######

            # OPTION 1: set "staticPasswords" as default auth type
            redirect_url_obj = redirect_url_obj._replace(
                path=re.sub(r"/auth$", "/auth/local", redirect_url_obj.path),
            )
            # OPTION 2: set "ldap" as default auth type
            # redirect_url_obj = redirect_url_obj._replace(
            #     path=re.sub(r"/auth$", "/auth/ldap", redirect_url_obj.path)
            # )

        # if we are at `/auth/xxxx/login` path, then no further action is needed (we can use it for login POST)
        if re.search(r"/auth/.*/login$", redirect_url_obj.path):
            auth_session["dex_login_url"] = redirect_url_obj.geturl()

        # else, we need to be redirected to the actual login page
        else:
            # this GET should redirect us to the `/auth/xxxx/login` path
            resp = s.get(redirect_url_obj.geturl(), allow_redirects=True)
            if resp.status_code != 200:
                raise RuntimeError(
                    f"HTTP status code '{resp.status_code}' for GET against: {redirect_url_obj.geturl()}",
                )

            # set the login url
            auth_session["dex_login_url"] = resp.url

        ################
        # Attempt Dex Login
        ################
        resp = s.post(
            auth_session["dex_login_url"],
            data={"login": username, "password": password},
            allow_redirects=True,
        )
        if len(resp.history) == 0:
            raise RuntimeError(
                f"Login credentials were probably invalid - "
                f"No redirect after POST to: {auth_session['dex_login_url']}",
            )

        # store the session cookies in a "key1=value1; key2=value2" string
        auth_session["session_cookie"] = "; ".join(
            [f"{c.name}={c.value}" for c in s.cookies]
        )

    return auth_session


def get_kubeflow_task_status(run_id: str, token: str) -> str:
    res = requests.get(
        f"{KUBEFLOW_API_V2}/{run_id}",
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
        },
    )
    try:
        # status = res.json()["run"]["status"] # v1api
        status = res.json()["state"]
    except Exception as e:
        logger.info(e)
        status = "Unknown"
    status = status[:1] + status[1:].lower()
    return status


def get_kubeflow_cookie() -> str:
    auth_session = get_istio_auth_session(
        KUBEFLOW_ENDPOINT, KUBEFLOW_USERNAME, KUBEFLOW_PASSWORD
    )
    return auth_session["session_cookie"]


def main():
    token = get_kubeflow_cookie()
    token = token.split("authservice_session=")[1]
    for gt_task in gt_task_col.find(
        {}, {"status": 1, "flow_run_id": 1, "created_at": 1}
    ):
        if gt_task.get("flow_run_id") is None:
            logger.warning(f"{gt_task['_id']} no flow_run_id")
            continue
        if gt_task.get("status") in ["Succeeded", "Failed"]:
            continue
        status = get_kubeflow_task_status(gt_task["flow_run_id"], token)
        if status == "Unknown":
            logger.info(f"task {gt_task['_id']} status is Unknown, skip")
            continue
        logger.info(f"update task {gt_task['_id']} status to {status}")
        gt_task_col.update_one(
            {"_id": gt_task["_id"]},
            {"$set": {"status": status}},
        )


if __name__ == "__main__":
    main()
